using System;
using System.Linq;
using System.IO;

using Microsoft.Cci;
using Microsoft.Cci.ILToCodeModel;

using ErrorEventArgs = Microsoft.Cci.ErrorEventArgs;

namespace ReactiveLang.AssemblyRewriter
{
    public class Rewriter
    {
        IMetadataHost _host = new PeReader.DefaultHost();

        public event Action<ErrorEventArgs> Errors = delegate { };

        public Rewriter()
        {
            _host.Errors += new EventHandler<ErrorEventArgs>(_host_Errors);
        }

        bool failed;
        void _host_Errors(object sender, ErrorEventArgs e)
        {
            if (e.Errors.Any(err => !err.IsWarning))
                failed = true;
            Errors(e);
        }

        public bool Rewrite(string inPe, string outPe)
        {
            failed = false;
            var module = _host.LoadUnitFrom(inPe) as IModule;
            if (module == null || module == Dummy.Module || module == Dummy.Assembly)
            {
                Console.WriteLine(inPe + " is not a PE file containing a CLR module or assembly.");
                return false;
            }
            string tmpPe = Path.GetTempFileName();
            try
            {
                string tmpPdb = Path.GetTempFileName();

                try
                {
                    //Get a PDB reader if there is a PDB file.
                    PdbReader pdbReader = null;
                    string pdbFile = Path.ChangeExtension(inPe, "pdb");
                    if (File.Exists(pdbFile))
                    {
                        Stream pdbStream = File.OpenRead(pdbFile);
                        pdbReader = new PdbReader(pdbStream, _host);
                    }

                    bool gotPdb = false;
                    using (pdbReader)
                    {
                        var mutator = new ReactiveCodeMutator(_host, pdbReader);
                        module = mutator.Visit(module);

                        if (!failed)
                        {
                            using (var peStream = File.Create(tmpPe))
                            {
                                if (pdbReader == null)
                                {
                                    PeWriter.WritePeToStream(module, _host, peStream);
                                }
                                else
                                {
                                    gotPdb = true;
                                    using (var pdbWriter = new PdbWriter(tmpPdb, pdbReader))
                                    {
                                        PeWriter.WritePeToStream(module, _host, peStream, pdbReader, pdbReader, pdbWriter);
                                    }
                                }
                            }
                        }
                    }


                    if (!failed)
                    {
                        File.Copy(tmpPe, outPe, true);
                        if (gotPdb)
                            File.Copy(tmpPdb, Path.ChangeExtension(outPe, ".pdb"), true);
                    }
                    return !failed;
                }
                finally
                {
                    File.Delete(tmpPdb);
                }
            }
            finally
            {
                File.Delete(tmpPe);
            }
        }
    }
}